package com.zzw.ratelimter;

import com.zzw.properties.RedisRateLimterProperties;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.RedisScript;

import java.util.Arrays;
import java.util.List;

/**
 * @author Administrator
 */
public class RedisRateLimiter implements RateLimiter {

    private static final Logger logger = LoggerFactory.getLogger(RedisRateLimiter.class);

    private final RedisRateLimterProperties redisRateLimterProperties;
    private final StringRedisTemplate redisTemplate;
    private final RedisScript<List<Long>> script;

    public RedisRateLimiter(StringRedisTemplate redisTemplate, RedisScript<List<Long>> script, RedisRateLimterProperties redisRateLimterProperties) {
        this.redisTemplate = redisTemplate;
        this.script = script;
        this.redisRateLimterProperties = redisRateLimterProperties;
    }


    static List<String> getKeys(String id) {
        String prefix = "request_rate_limiter.{" + id;
        String tokenKey = prefix + "}.tokens";
        String timestampKey = prefix + "}.timestamp";
        return Arrays.asList(tokenKey, timestampKey);
    }

    public Response isAllowed(String id) {
        List<String> keys = getKeys(id);
        List<Long> exec = null;
        try {
            exec = this.redisTemplate.execute(this.script, keys, String.valueOf(this.redisRateLimterProperties.getReplenishRate())
                    , String.valueOf(this.redisRateLimterProperties.getBurstCapacity())
                    , ""
                    , String.valueOf(this.redisRateLimterProperties.getRequestedTokens()));
        } catch (Throwable throwable) {
            this.logger.error("Error calling rate limiter lua", throwable);
            exec = Arrays.asList(1L, -1L);
        }
        boolean allowed = exec.get(0) == 1L;
        Long tokensLeft = exec.get(1);
        return new Response(allowed, tokensLeft);
    }

    public Response isAllowed(String id, int qps) {
        List<String> keys = getKeys(id);
        List<Long> exec = null;
        try {
            exec = this.redisTemplate.execute(this.script, keys, qps + ""
                    , qps + ""
                    , ""
                    , "1");
        } catch (Throwable throwable) {
            this.logger.error("Error calling rate limiter lua", throwable);
            exec = Arrays.asList(1L, -1L);
        }
        boolean allowed = exec.get(0) == 1L;
        Long tokensLeft = exec.get(1);
        return new Response(allowed, tokensLeft);
    }
}
